# 本脚本用于测试无kv-server的情况下，CPU顺序执行的吞吐情况（无pipeline）

import subprocess
import os
import csv
import math

# 创建结果目录
results_dir = '../results_raw'
os.makedirs(results_dir, exist_ok=True)

# 结果文件
output_file = os.path.join(results_dir, 'sequence_throughput.csv')

# 定义批大小范围（从1024到524288，每次翻倍）
batch_sizes = [2**i for i in range(1, 20)]  # 1024 (2^10) 到 524288 (2^19)

def convert_to_us(value, unit):
    """将时间单位转换为微秒"""
    if unit == 's':
        return value * 1_000_000
    elif unit == 'ms':
        return value * 1_000
    elif unit == 'us':
        return value
    elif unit == 'ns':
        return value / 1000
    else:
        raise ValueError(f"Unknown unit: {unit}")

# 打开文件以写入结果
with open(output_file, 'w', newline='') as csvfile:
    fieldnames = [
        'batch_size',
        'p50_time_us',
        'p99_time_us',
        'data_throughput_MB/s_p50',
        'data_throughput_MB/s_p99',
        'request_throughput_rps_p50',
        'request_throughput_rps_p99'
    ]
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    writer.writeheader()

    for batch_size in batch_sizes:
        # 提示当前运行的批大小
        print(f"Running test with batch_size={batch_size}")
        
        # 构建命令
        command = [
            './../../bin_host/test_cpu',
            f'--batch_size={batch_size}',
            '--if_pipeline=0'
        ]

        # 执行命令并捕获输出
        try:
            result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, 
                                    text=True, encoding='utf-8', errors='ignore')
        except Exception as e:
            print(f"Error executing command: {e}")
            continue

        stdout = result.stdout
        
        # 提取包含"total_time"的行
        time_line = None
        for line in stdout.splitlines():
            if "total_time" in line:
                time_line = line
                break
        
        if not time_line:
            print(f"Warning: No 'total_time' found for batch_size={batch_size}")
            continue
        
        # 解析表格行（格式：| total_time | 92.484 ms | 127.508 ms |）
        parts = time_line.split('|')
        if len(parts) < 4:
            print(f"Warning: Unexpected format for batch_size={batch_size}")
            continue
        
        # 提取P50和P99的值和单位
        p50_str = parts[2].strip()
        p99_str = parts[3].strip()
        
        # 分离数值和单位（例如："92.484 ms" -> 92.484 和 "ms"）
        p50_value, p50_unit = p50_str.split()
        p99_value, p99_unit = p99_str.split()
        
        # 转换为微秒
        p50_us = convert_to_us(float(p50_value), p50_unit)
        p99_us = convert_to_us(float(p99_value), p99_unit)
        
        # 计算数据大小（每个batch_size有40个键，每个键8字节）
        data_bytes = batch_size * 40 * 8
        
        # 计算吞吐量（数据吞吐量：MB/s）
        data_throughput_p50 = data_bytes / (p50_us / 1_000_000) / (1024 * 1024)
        data_throughput_p99 = data_bytes / (p99_us / 1_000_000) / (1024 * 1024)
        
        # 计算请求吞吐量（requests per second）
        request_throughput_p50 = batch_size / (p50_us / 1_000_000)
        request_throughput_p99 = batch_size / (p99_us / 1_000_000)
        
        # 写入CSV行
        writer.writerow({
            'batch_size': batch_size,
            'p50_time_us': p50_us,
            'p99_time_us': p99_us,
            'data_throughput_MB/s_p50': data_throughput_p50,
            'data_throughput_MB/s_p99': data_throughput_p99,
            'request_throughput_rps_p50': request_throughput_p50,
            'request_throughput_rps_p99': request_throughput_p99
        })

print(f"测试完成! 结果已保存到 {output_file}")
